{-# LANGUAGE ViewPatterns, DeriveDataTypeable #-}
----------------------------------------------------------------
--                                                            --
-- Module      :  LLVM.Code.Function                          --
-- License     :  GPL3                                        --
-- Mantainer   :  Mikhail Belyaev                             --
--                                                            --
----------------------------------------------------------------

module LLVM.Code.Function where

import Data.Generics(Typeable, Data)

import LLVM.Parser.Record.Record
       (extrudeCustomSubblocksById, recordAsHeadlessList, getBlockId,
        RawBlockElem(..), RawBlockElem, fromData, fromSub, BlockHeader(..),
        BlockType(..), extrudeCustomSubblocks, extrudeCustomRecords,
        recordAsList, RawBlock(..), RecordType(..))
import LLVM.Code.Shared
       (BBIndex(..), ValIndex(..), VisibilityType(..), LinkageType(..),
        CallingConv(..), ParamAttrIndex(..),
        SectionIndex(..), GCIndex(..), VisibilityType, LinkageType,
        CallingConv, TypeIndex(..))
import LLVM.Code.ParamAttrs (ParamAttributes(..))
import LLVM.Util
import LLVM.Code.Instructions.Instruction
       (isInstruction, isTerminating)
import LLVM.Code.Constant (emptyConsts, ConstantBlock(..))
import LLVM.Code.Symbolic (emptyVS, ValSymbolTable(..))
import Data.List
       (findIndices, genericLength, genericReplicate, find, mapAccumL)
import Control.Monad (liftM)
import qualified Data.DataMap as M (assocs, lookup, fromList, Map, adjust )
import LLVM.Code.Values
       (getValue, vtSize, vtFromList, ValueTable(..), TypedValue(..),
        Value(..), Value)
import LLVM.Code.Values.Kinds
       (InstructionValue(..), ParameterValue(..), FunctionAttrs(..),
        FunctionDeclValue(..), MetaValue(..))
import LLVM.Code.Types
       (Type(..), getTypeByIndex, Type, TypeTable(..))
import Data.Maybe (fromJust, catMaybes)
import LLVM.Code.DebugLoc (DebugLocAgain(..), DebugTable(..))
import LLVM.Code.Instructions.Class (BuildableInstruction(..))
import Data.Monoid (Monoid(..))
import Data.DataMap (findMax)
import LLVM.Code.Instructions.Values
       (StoreInstruction(..), LoadInstruction(..), FreeInstruction(..),
        AllocaInstruction(..), MallocInstruction(..),
        InvokeInstruction(..), CallInstruction(..), SwitchInstruction(..),
        IndirBranchInstruction(..), BranchInstruction(..),
        InsertValueInstruction(..), ExtractValueInstruction(..),
        PhiInstruction(..), ReturnInstruction(..), CmpInstruction(..),
        ShuffleVecInstruction(..), InsertElementInstruction(..),
        ExtractElementInstruction(..), SelectInstruction(..),
        GetElementPtrInstruction(..), CastInstruction(..),
        BinaryInstruction(..), Instruction(..), DebugLoc(..), Instruction)

data FunctionDeclDef = FunctionDeclDef {
                        functionDDType :: TypeIndex,
                        functionDDValue :: FunctionDeclValue
                      } deriving (Show,Eq,Typeable,Data)

instance RecordType FunctionDeclDef where
    fromRawRecord
        (recordAsHeadlessList ->
            8:                              -- function decl
            (TypeIndex -> tp):               -- type
            (integerToIC -> Just cc):
            ((/=0) -> isproto):
            (integerToIC -> Just ln):
            pa:
            al:
            sec:
            (integerToIC -> Just vis):
            gc) =
                                Just $ FunctionDeclDef{
                                    functionDDType = tp,
                                    functionDDValue = FunctionDeclValue {
                                        isAbstract = isproto,
                                        functionDeclAttrs =
                                            FunctionAttrs cc ln rpa al rsec vis rgc
                                    }
                                }
                                where nullOr1based i = if(i/=0) then Just $ fromInteger (i-1)
                                                               else Nothing
                                      rpa = nullOr1based pa
                                      rsec = nullOr1based sec
                                      rgc = case gc of
                                                 [] -> Nothing
                                                 (x:_) -> nullOr1based x
    fromRawRecord _ = Nothing

instance TypedValue FunctionDeclDef where
    getType (functionDDType -> tp) tt = tt `getTypeByIndex` tp ?? LLOpaque
    toValue (functionDDValue -> v) = FunctionV v


type BasicBlockTable = M.Map BBIndex BasicBlock
data BasicBlock = BasicBlock {
                        bbFirst :: ValIndex,
                        bbLast :: ValIndex,
                        bbPreceeding :: [BBIndex],
                        bbFollowing :: [BBIndex]
                  } deriving (Show,Eq,Typeable,Data)

bbFromValueTable :: ValueTable -> BasicBlockTable
bbFromValueTable vt@(ValueTable mp) =
                      let binds = M.assocs mp
                          isInstr (_, (_, InstructionV _)) = True
                          isInstr _ = False
                          funcVT = snd $ break isInstr binds
                          isTerm (_, (_, InstructionV (InstructionValue instr _))) =
                            isTerminating instr
                          isTerm _ = False
                          bbs = groups isTerm funcVT
                      in makeLinks vt $ zip [0..] $
                        map (\ vals -> BasicBlock (fst $ head vals) (fst $ last vals) [] []) bbs

makeLinks :: ValueTable -> [(BBIndex,BasicBlock)] -> BasicBlockTable
makeLinks vt prs = addBackwards $ addForwards prs
    where forward bb = case (vt `getValue` bbLast bb) of
                            (Just (InstructionV (InstructionValue inst _))) ->
                                case inst of
                                     (IRet _) -> []
                                     (IBranch (StraightBranch bbi)) -> [bbi]
                                     (IBranch (CondBranch bb1 bb2 _)) -> [bb1,bb2]
                                     (IIndirBr (IndirBranchInstruction{ibrPossibleValues = bbs})) ->
                                        bbs
                                     (IInvoke (InvokeInstruction{normalBB = bb1, catchBB = bb2})) ->
                                        [bb1,bb2]
                                     (IUnwind _) -> []
                                     (IUnreachable _) -> []
                            _ -> []
          addForwards ((i,bb):bbs) = (i,bb{bbFollowing = forward bb}) : addForwards bbs
          addForwards [] = []
          backLinks ((i,bb):bbs) = map (\x -> (i,x)) (forward bb) ++ backLinks bbs
          backLinks [] = []
          addbLink lnk bb = bb{bbPreceeding = lnk : bbPreceeding bb}
          addbLinks bbs (x,y) = M.adjust (addbLink x) y bbs
          addBackwards bbs = foldl addbLinks (M.fromList bbs) (backLinks bbs)

-- (IRet _ ) = True
-- (IBranch _) = True
-- (ISwitch _) = True
-- (IIndirBr _) = True
-- (IInvoke _) = True
-- (IUnwind _) = True
-- (IUnreachable _) = True

data FunctionBodyDef = FunctionBodyDef {
                            funcLocalValSyms :: ValSymbolTable,
                            funcLocalVals :: ValueTable,
                            funcBasicBlocks :: BasicBlockTable
                       } deriving (Show,Eq,Typeable,Data)

newtype FunctionTable = FunctionTable { bodies :: [FunctionBodyDef] } deriving (Show,Eq,Typeable,Data)

addInterm :: ValueTable -> [(Type,Value)] -> ValueTable
addInterm (ValueTable mp) assocs = ValueTable $ mp `mappend` umap
    where start = fst $ findMax mp
          zipped = zip [start+1..] assocs
          umap = M.fromList $ catMaybes $ map Just zipped

collectFunctions :: RawBlock -> TypeTable -> ValueTable -> FunctionTable
collectFunctions (contents -> conts) tt vt =
                        let isFDef (_,(_,FunctionV (isAbstract -> False))) = True
                            isFDef _ = False
                            defs = map (fst.snd) $ case vt of
                                        ValueTable mp -> filter isFDef $ M.assocs mp
                            bods = filter (\x -> getBlockId x == 12) $ maybeFilter fromSub conts
                            zipped = zip defs bods
                            extractParsTo (LLPointer ti _) vt =
                                let pars = case (tt `getTypeByIndex` ti) of
                                                (Just (LLFunctionType _ prs _)) -> prs
                                                _ -> []
                                in
                                vt `addInterm`
                                zipWith
                                    (\x y -> (fromJust $ tt `getTypeByIndex` x,ParameterV $ ParameterValue y))
                                    pars [0..]
                            extractParsTo _ vt = vt
                            ept t vt = extractParsTo t vt
                        in
                        FunctionTable {
                            bodies = maybeFilter (\(t,x) -> buildFunctionBody x tt $ ept t vt) zipped
                        }

buildFunctionBody :: RawBlock -> TypeTable -> ValueTable -> Maybe FunctionBodyDef
buildFunctionBody rb@(getBlockId -> 12) tt vt =
                       Just FunctionBodyDef {
                            funcLocalValSyms = valSym,
                            funcLocalVals = newVT,
                            funcBasicBlocks = bbFromValueTable newVT
                       } where
                            valSym = (mayhead $ rb `extrudeCustomSubblocksById` 14) ?? emptyVS
                            newVT = readFuncValueTable tt vt rb
buildFunctionBody _ _ _ = Nothing


groups :: (a -> Bool) -> [a] -> [[a]]
groups f l = case (break f l) of
                  (bef,m:[]) -> [bef++[m]]
                  (bef,m:af) -> [bef++[m]] ++ groups f af
                  (bef,[]) -> [bef]

--collectInstructions :: RawBlock -> [Instruction]
--collectInstructions rb = fromRawRecs recs
--                         where recs = maybeFilter fromData $ rb.$contents
--
--collectBasicBlocks :: RawBlock -> [BasicBlockDef]
--collectBasicBlocks rb = map BasicBlockDef instrGroups
--                        where instrGroups = groups isTerminating $ collectInstructions rb
--
--instance BlockType FunctionBodyDef where
--    fromRawBlock rb@( blockId . header -> 12 ) = Just FunctionBodyDef {
--                            funcLocalConstants = consts,
--                            funcLocalValSyms = valSym,
--                            funcLocalVals = readFuncValueTable undefined rb 0,
--                            basicBlocks = rb.$collectBasicBlocks
--                       }
--                       where
--                            conts = rb.$contents
--                            subs = maybeFilter fromSub conts
--                            blockById v = find ( (==v) . blockId . header ) subs
--                            vsBlock = blockById 14
--                            valSym = (flop $ liftM fromRawBlock vsBlock) ?? emptyVS
--                            csBlock = blockById 11
--                            consts = (flop $ liftM fromRawBlock csBlock) ?? emptyConsts
--
--    fromRawBlock _ = Nothing

isNotVoid :: (Type, Value) -> Bool
isNotVoid (LLVoid, _) = False
isNotVoid (_,_) = True

isVoid :: (Type, Value) -> Bool
isVoid x = not $ isNotVoid x

-- this function is double-reversed!!! (reversed list on input, reversed list on output)
modValueTable :: TypeTable -> ValueTable -> RawBlockElem -> [(Type,Value)] -> [(Type,Value)]
modValueTable tt vt (Data rr@(isInstruction -> True)) xs =
    case (buildInstruction rr tt $ vt `addInterm` filter isNotVoid xs)::Maybe Instruction of
         Just instr -> xs ++ [toTVPair instr tt]
         _ -> xs

modValueTable tt vt (Data (fromRawRecord -> Just dl@DebugLoc{})) xs =
    case (last xs) of
         (t, InstructionV (InstructionValue inst Nothing)) ->
            init xs ++ [(t, InstructionV (InstructionValue inst $ Just dl))]
         _ -> xs

modValueTable tt vt (Data (fromRawRecord -> Just dl@DebugLocAgain{})) xs =
    let fetchDL (_,InstructionV (InstructionValue _ dl)) = dl
        fetchDL _ = Nothing
        lastDL = head $ maybeFilter fetchDL $ reverse xs
    in
    case (last xs) of
         (t, InstructionV (InstructionValue inst Nothing)) ->
            init xs ++ [(t, InstructionV (InstructionValue inst $ Just lastDL))]
         _ -> xs

-- subblock with id 11 is constant block
modValueTable
        tt
        _ -- value table isn't needed
        (Sub (mmatch -> (            -- now that's what I call a pattern match )))
                getBlockId -> 11,
                fromRawBlock -> Just (cdefs -> cdefs)
        )))
        xs
            = xs ++ map (`toTVPair` tt) cdefs
modValueTable _ _ _ x = x

readFuncValueTable :: TypeTable -> ValueTable -> RawBlock -> ValueTable
readFuncValueTable tt vt rb =
                          vtFromList ret
                          -- note using foldr instead of foldl (implicit reverse)
                          where valkinds = foldl (flip $ modValueTable tt vt) [] $ rb.$contents
                                start = vtSize vt
                                shifted = tiltShift start $ putMetas valkinds
                                zipped = zip [start..] shifted
                                ret = catMaybes $ map Just zipped

tiltShift :: ValIndex -> [(Type,Value)] -> [(Type,Value)]
tiltShift start lp = let
                        voidIxs = map (\ x -> toInteger x + toInteger start) $ findIndices isVoid lp
                        shft (t,x) y = (t,valShift x 1 y)
                        ashft lst y = map (`shft` y) lst
                     in
                     foldl ashft lp voidIxs


tiltShifte :: [(Type,Value)] -> Integer -> [(Type,Value)]
tiltShifte ((LLVoid,v):rest) i = (LLVoid,valShift v i i):(tiltShifte rest $ i+1)
tiltShifte ((t,v):rest) i = (t,valShift v i i):(tiltShifte rest i)
tiltShifte [] _ = []

isCall :: (Type,Value) -> Bool
isCall (_,InstructionV (InstructionValue (ICall _) _)) = True
isCall (_,InstructionV (InstructionValue (IInvoke _) _)) = True
isCall _ = False

callParams :: (Type,Value) -> [ValIndex]
callParams (_,InstructionV (InstructionValue (ICall (callOperands -> ops)) _)) = ops
callParams (_,InstructionV (InstructionValue (IInvoke (nestedCall -> (callOperands -> ops))) _)) = ops
callParams _ = []

putMetas :: [(Type,Value)] -> [(Type,Value)]
putMetas (cl@(callParams -> pars):rest) =
    let
        metas = map (\x -> (LLVoid, MetaV $ MetaAccessor x)) pars
    in
    cl:metas ++ putMetas rest
putMetas [] = []

valShift :: Value -> Integer -> Integer -> Value
valShift (InstructionV (InstructionValue instr dl)) i min =
    InstructionV $ InstructionValue (valShiftI instr min i) dl
valShift (MetaV (MetaAccessor tgt)) i min =
    let mod vi = if(toInteger vi >= min) then fromInteger $ toInteger vi + i
                                             else vi
    in
    MetaV (MetaAccessor (mod tgt))
valShift v _ _ = v

valShiftI :: Instruction -> Integer -> Integer-> Instruction
valShiftI inst min i =
        let viPlus vi i = if(toInteger vi >= min) then fromInteger $ toInteger vi + i
                                                 else vi
            re :: ValIndex -> ValIndex
            re vi = vi `viPlus` i
        in case inst of
                (IExtractEl is@ExtractElementInstruction{
                    extractElIx = (re -> ix), extractElVal = (re -> v)
                 }) -> IExtractEl is{extractElIx = ix, extractElVal = v}
                (IInsertEl is@InsertElementInstruction{
                    insertElIx = (re -> ix), insertElVal = (re -> v), insertElElem = (re -> e)
                }) -> IInsertEl is{insertElIx = ix, insertElVal = v, insertElElem = e}
                (IShuffleVec is@ShuffleVecInstruction{
                    shuffleVec1 = (re -> v1), shuffleVec2 = (re -> v2), shuffleMask = (re -> m)
                }) -> IShuffleVec is{shuffleVec1 = v1, shuffleVec2 = v2, shuffleMask = m}
                (IExtractV is@ExtractValueInstruction{ extractValVal = (re -> v) }) ->
                    IExtractV is{ extractValVal = v}
                (IInsertV is@InsertValueInstruction{
                    insertValVal = (re -> v),
                    insertValElem = (re -> el)
                }) -> IInsertV is { insertValVal = v, insertValElem = el }
                (IGEP gep@GetElementPtrInstruction{gepVal = (re -> val), gepIxs = ixs}) ->
                    IGEP gep{gepVal = val, gepIxs = map (\(x,y) -> (x,y `viPlus`i)) ixs}
                (IBinary bi@BinaryInstruction{biOp1 = (re -> op1), biOp2 = (re -> op2)}) ->
                    IBinary bi{biOp1 = op1, biOp2 = op2}
                -- branch instr (only conditional kind has values in it)
                (IBranch (CondBranch bb0 bb1 (re -> v) )) ->
                    IBranch $ CondBranch bb0 bb1 v
                (IIndirBr is@IndirBranchInstruction{ibrVal = (re -> v) }) ->
                    IIndirBr is{ ibrVal = v }
                (ISwitch is@SwitchInstruction{swOpval = (re -> v), swCases = cases}) ->
                    ISwitch is{swOpval = v, swCases = map (\((re -> v),bb) -> (v,bb)) cases}
                (ICall is@CallInstruction{callFValue = (re -> v), callOperands = (map re -> ops) }) ->
                    ICall is{callFValue = v, callOperands = ops}
                (IInvoke is@InvokeInstruction{
                    nestedCall = ic@CallInstruction{callFValue = (re -> v), callOperands = (map re -> ops) }
                }) -> IInvoke is{nestedCall = ic{callFValue = v, callOperands = ops}}
                -- return need shifting only if it has a value
                (IRet is@ReturnInstruction{ result = Just (t,(re -> v)) }) ->
                    IRet is{result = Just(t,v)}
                -- Unwind has no values
                -- Unreachable has no values
                (ICast ci@CastInstruction{ciOp = (re -> op)}) ->
                    ICast ci{ciOp = op}
                (ICmp is@CmpInstruction{ cmpOp1Val = (re -> op1), cmpOp2Val = (re -> op2) }) ->
                    ICmp is{cmpOp1Val = op1, cmpOp2Val = op2}
                (ISelect is@SelectInstruction{
                    selOpTrue = (re -> tr), selOpFalse = (re -> fls), selPred = (re -> pr)
                 }) -> ISelect is{selOpTrue = tr, selOpFalse = fls, selPred = pr}
                (IPhi is@PhiInstruction{ phiVals = phiv } )->
                    IPhi is{phiVals = map (\((re -> v),bb) -> (v,bb)) phiv}
                (IMalloc is@MallocInstruction { mallocOp = (re -> op) }) ->
                    IMalloc is{ mallocOp = op }
                (IAlloca is@AllocaInstruction { allocaOp = (re -> op) }) ->
                    IAlloca is{ allocaOp = op }
                (IFree is@FreeInstruction { freeValue = (re -> v) }) ->
                    IFree is{ freeValue = v }
                (ILoad is@LoadInstruction { loadPtr = (re -> ptr) }) ->
                    ILoad is{ loadPtr = ptr }
                (IStore is@StoreInstruction { storePtr = (re -> ptr), storeVal = (re -> val) }) ->
                    IStore is{ storePtr = ptr, storeVal = val }
                some -> some



